import os
import argparse
import paddle


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', type=str)
    parser.add_argument('--output', type=str)
    parser.add_argument('--type', type=int, default=0)
    args = parser.parse_args()
    return args


def is_valid_keys(name, keys):
    for k in keys:
        if name.startswith(k):
            return True
    return False
    

def main():
    args = parse_args()
    sd = paddle.load(args.input)
    pop_keys = []
    if args.type == 0:
        # keep_keys = ['backbone.', 'neck.']
        keep_keys = ['backbone.']
    elif args.type == 1:
        keep_keys = ['backbone.', 'neck.', 
            'head.binarize.conv1', 
            'head.binarize.conv_bn1',
            'head.binarize.conv2',
            'head.binarize.conv_bn2',
            'head.thresh.conv1', 
            'head.thresh.conv_bn1',
            'head.thresh.conv2',
            'head.thresh.conv_bn2']

    for k in sd.keys():
        if not is_valid_keys(k, keep_keys):
            pop_keys.append(k)
            print(k)
    for k in pop_keys:
        sd.pop(k)
    paddle.save(sd, args.output)


if __name__ == '__main__':
    main()
